import asyncio
from typing import Dict, Optional,Any
from fastapiApp.utils import udp_socket
from fastapiApp.utils import simulation_client
from enum import Enum

# 任务状态枚举
class TaskStatus(str, Enum):
    RUNNING = "running"
    STOPPED = "stopped"
    ERROR = "error"
    DONE = "done"

# 任务信息模型，包含更多详细信息
class TaskInfo:
    def __init__(self):
        self.task: Optional[asyncio.Task] = None
        self.status: TaskStatus = TaskStatus.STOPPED
        self.error: Optional[str] = None  # 存储错误信息
        self.result: Optional[Any] = None  # 存储任务返回结果
        self.description: Optional[str] = None  # 备注

# 任务管理模型
class TaskManager:
    def __init__(self):
        self.tasks: Dict[str, TaskInfo] = {
            "simulation_client": TaskInfo(),
            "udp_socket": TaskInfo()
        }
        # 任务函数映射
        self.task_functions = {
            "simulation_client": simulation_client.async_simulation_clients,
            "udp_socket": udp_socket.async_udp_clients
        }

        self.tasks["simulation_client"].description = "模拟设备进程"
        self.tasks["udp_socket"].description = "UDP通信进程"


    async def start_task(self, task_name: str):
        task_info = self.tasks[task_name]

        if task_info.status == TaskStatus.RUNNING:
            return False

        # 重置任务信息
        task_info.error = None
        task_info.result = None

        # 包装任务函数以捕获异常
        async def task_wrapper():
            try:
                # 执行实际任务
                result = await self.task_functions[task_name]()
                return result
            except Exception as e:
                # 捕获所有异常并存储
                task_info.error = str(e)
                raise  # 重新抛出以便外部处理

        # 创建新任务
        task = asyncio.create_task(
            task_wrapper(),
            name=task_name
        )

        # 任务完成时更新状态
        def task_done_callback(fut):
            try:
                # 检查任务是否有结果
                result = fut.result()
                task_info.result = result
                task_info.status = TaskStatus.DONE
            except asyncio.CancelledError:
                # 任务被取消
                task_info.status = TaskStatus.STOPPED
            except Exception:
                # 任务抛出异常
                task_info.status = TaskStatus.ERROR
            finally:
                task_info.task = None

        task.add_done_callback(task_done_callback)
        task_info.task = task
        task_info.status = TaskStatus.RUNNING
        return True

    async def stop_task(self, task_name: str):
        task_info = self.tasks[task_name]

        if task_info.status != TaskStatus.RUNNING:
            return False

        # 取消任务
        if task_info.task:
            task_info.task.cancel()
            try:
                await task_info.task
            except asyncio.CancelledError:
                pass

        task_info.tasks = None
        task_info.status = TaskStatus.STOPPED
        return True

    async def restart_task(self, task_name: str):
        # 先停止再启动
        await self.stop_task(task_name)
        return await self.start_task(task_name)

    def get_task_details(self, task_name: str):
        """获取任务的详细信息"""
        task_info = self.tasks[task_name]
        return {
            "status": task_info.status,
            "error": task_info.error,
            "result": task_info.result,
            "is_running": task_info.status == TaskStatus.RUNNING,
            "description": task_info.description,
        }

# 创建任务管理器实例
task_manager = TaskManager()
